'''
Created on Jul 14, 2012

@author: nzhao
'''
import numpy as np

from traits.api import HasTraits, String, Instance, Enum, Array
from traitsui.api import View, Item, Group, ArrayEditor
from misc.dynamicRange import DynamicRangeFloat

def pol_group(objname):
    g = Group(Item( objname+'.name', show_label=True, style='readonly'),
                Group(
                      Item(objname+'.magnitude', style= 'custom', label = 'Magnitude', show_label = True),
                      Item(objname+'.theta',     style= 'custom', label = 'Theta    ', show_label = True),
                      Item(objname+'.phi',       style= 'custom', label = 'Phi      ', show_label = True),
                      Item(objname+'.xyz',       style= 'readonly', show_label = False, 
                           editor = ArrayEditor(width = 150, format_str = '%8.3f')),
                      )
                )
    return g
def xyz_group(objname):
    g = Group(
                    Item(objname+'.name', show_label=True, style='readonly'),
                    Item(objname+'.x', style= 'custom', show_label = False),
                    Item(objname+'.y', style= 'custom', show_label = False),
                    Item(objname+'.z', style= 'custom', show_label = False),
                    Item(objname+'.pol', style= 'readonly', show_label = False, 
                         editor = ArrayEditor(width = 150, format_str = '%8.3f')),
                    )
    return g

class vector(DynamicRangeFloat):
    name = String()
    type = Enum('Cartesian', 'Sphere')
    magnitude = Instance(DynamicRangeFloat, ())
    theta     = Instance(DynamicRangeFloat, ())
    phi       = Instance(DynamicRangeFloat, ())
    x         = Instance(DynamicRangeFloat, ())
    y         = Instance(DynamicRangeFloat, ())
    z         = Instance(DynamicRangeFloat, ())
    xyz       = Array()
    pol       = Array()
    
    def __init__(self, xyz = [0., 0., 0.], pol = None, 
                 name = 'Vector', type = 'Cartesian',
                 magnitude_range = (0.0, np.inf), theta_range = (0.0, 180.0),  phi_range = (0.0, 360.0),
                 x_range = (-np.inf, np.inf),     y_range = (-np.inf, np.inf), z_range = (-np.inf, np.inf) 
                 ):
        self.name = name
        self.type = type
        if pol is not None:
            self.type = 'Sphere'
        
        self.x = DynamicRangeFloat(name, min_val = x_range[0], max_val = x_range[1])
        self.y = DynamicRangeFloat(name, min_val = y_range[0], max_val = y_range[1])
        self.z = DynamicRangeFloat(name, min_val = z_range[0], max_val = z_range[1])
        
        self.magnitude = DynamicRangeFloat(name, min_val = magnitude_range[0], max_val = magnitude_range[1])
        self.theta     = DynamicRangeFloat(name, min_val = theta_range[0],     max_val = theta_range[1])
        self.phi       = DynamicRangeFloat(name, min_val = phi_range[0],       max_val = phi_range[1])

        if self.type is 'Cartesian':
            self.x.value, self.y.value, self.z.value = xyz
            self.xyz = np.array(xyz)
            self.xyz_2_pol()
            self.x.on_trait_change(self.xyz_2_pol, 'value')
            self.y.on_trait_change(self.xyz_2_pol, 'value')
            self.z.on_trait_change(self.xyz_2_pol, 'value')        
        elif self.type is 'Sphere':
            self.magnitude.value, self.theta.value, self.phi.value = pol
            self.pol = np.array(pol)
            self.pol_2_xyz()
            self.magnitude.on_trait_change(self.pol_2_xyz, 'value')
            self.theta.on_trait_change    (self.pol_2_xyz, 'value')
            self.phi.on_trait_change      (self.pol_2_xyz, 'value')
        else:
            print 'error'
    
    
    def set_val(self, val, type=None):
        if (type is 'Cartesian') or (type is 'Sphere'):
            self.type = type
             
        if self.type is 'Cartesian':
            self.xyz = val
            self.x.value, self.y.value, self.z.value = val
            self.xyz_2_pol()
        elif self.type is 'Sphere':
            self.pol = val
            self.magnitude.value, self.theta.value, self.phi.value = val
            self.pol_2_xyz()
        
    def _type_changed(self):
        if self.type == 'Cartesian':
            self.x.on_trait_change(handler = self.xyz_2_pol, name = 'value')
            self.y.on_trait_change(handler = self.xyz_2_pol, name = 'value')
            self.z.on_trait_change(handler = self.xyz_2_pol, name = 'value')
            self.magnitude.on_trait_change(handler = self.pol_2_xyz, name = 'value', remove = True)
            self.theta.on_trait_change    (handler = self.pol_2_xyz, name = 'value', remove = True)
            self.phi.on_trait_change      (handler = self.pol_2_xyz, name = 'value', remove = True)
        elif self.type == 'Sphere':
            self.x.on_trait_change(handler = self.xyz_2_pol, name = 'value', remove = True)
            self.y.on_trait_change(handler = self.xyz_2_pol, name = 'value', remove = True)
            self.z.on_trait_change(handler = self.xyz_2_pol, name = 'value', remove = True)
            self.magnitude.on_trait_change(handler = self.pol_2_xyz, name = 'value')
            self.theta.on_trait_change    (handler = self.pol_2_xyz, name = 'value')
            self.phi.on_trait_change      (handler = self.pol_2_xyz, name = 'value')
        else:
            print 'error type'                        
        
    def xyz_2_pol(self):
        x, y, z = self.x.value, self.y.value, self.z.value
        mag   = np.sqrt(x**2 + y**2 + z**2)
        if mag > 0.0:
            theta = np.arccos(z/mag) / np.pi * 180.0
            if theta == 0.0:
                phi = 0.0
            elif x == 0.0 and y >= 0.0 and theta != 0.0:
                phi = 90.0
            elif x == 0.0 and y < 0.0 and theta != 0.0:
                phi = -90.0
            else:
                phi = np.arctan(y/x) / np.pi * 180.0
        else:
            theta = 0.0
            phi   = 0.0
        
        if x < 0.0:
            phi = phi + 180.0
        
        self.magnitude.value, self.theta.value, self.phi.value = mag, theta, phi
        self.pol = np.array([mag, theta, phi])
        return self.pol
    
    def pol_2_xyz(self):
        mag, theta, phi = self.magnitude.value, self.theta.value / 180.0 * np.pi, self.phi.value / 180.0 * np.pi
        x = mag * np.sin(theta) * np.cos(phi)
        y = mag * np.sin(theta) * np.sin(phi)
        z = mag * np.cos(theta)
        self.x.value, self.y.value, self.z.value = x, y, z
        self.xyz = np.array([x,y,z])
        return self.xyz

class Cartesian_Vector(vector):
    def __init__(self, xyz = [0., 0., 0.],
                 name = 'Vector',  x_range = (-np.inf, np.inf),     y_range = (-np.inf, np.inf), z_range = (-np.inf, np.inf) ):
        
        super(Cartesian_Vector, self).__init__(name = name, type = 'Cartesian',
                                               xyz = xyz, x_range = x_range, y_range = y_range, z_range = z_range
                                               )
    traits_view=View(Group(Item('name', show_label=True, style='readonly'),
                    Item('x', style= 'custom', show_label = False),
                    Item('y', style= 'custom', show_label = False),
                    Item('z', style= 'custom', show_label = False),
                    Item('pol', style= 'readonly', show_label = False,
                        editor = ArrayEditor(width = 150, format_str = '%8.3f')),
                    )
                ) 

class Sphere_Vector(vector):
    def __init__(self, pol = [0., 0., 0.], 
                 name = 'Vector', magnitude_range = (0.0, np.inf), theta_range = (0.0, 180.0),  phi_range = (0.0, 360.0)):
        
        super(Sphere_Vector, self).__init__(name = name, type = 'Sphere',
                                               pol = pol, magnitude_range = magnitude_range, 
                                               theta_range = theta_range,  phi_range = phi_range
                                               )
    trait_view=View(Item('name', show_label=True, style='readonly'),
                    Group(
                          Item('magnitude', style= 'custom', label = 'Magnitude', show_label = True),
                          Item('theta',     style= 'custom', label = 'Theta    ', show_label = True),
                          Item('phi',       style= 'custom', label = 'Phi      ', show_label = True),
                          Item('xyz',       style= 'readonly', show_label = False, 
                           editor = ArrayEditor(width = 150, format_str = '%8.3f')),
                          )
                    )       

if __name__=='__main__':
    class test(HasTraits):
        xx = String
        bxyz=Instance(Cartesian_Vector, ())
        bpol=Instance(Sphere_Vector, ())
        
        def __init__(self):
            bxyz=Cartesian_Vector()
            bpol=Sphere_Vector()
        
        traits_view = View(
                           Item('xx'),
                           Group(Item('bxyz', style='custom' ),
                                 Item('bpol', style='custom' ),
                                 ),
                           width     = .8,
                           height    = .9,
                    )
    t = test()
    t.configure_traits()
    